# 采集笔记数据模块

import csv
import os
import time
from datetime import datetime

class DataCollector:
    def __init__(self):
        # 确保data目录存在
        if not os.path.exists('data/notes'):
            os.makedirs('data/notes')
        
        # 使用固定的CSV文件名
        self.csv_file = 'data/notes/xhs_python_notes.csv'
        
        # 定义CSV表头
        self.headers = [
            'title',          # 文章标题
            'author',         # 作者
            'content',        # 文章内容
            'edit_time',      # 编辑时间
            'likes',          # 点赞量
            'collects',       # 收藏量
            'comments',       # 评论量
            'collect_time',   # 采集时间
            'url',            # 文章URL
            'keyword'         # 搜索关键词
        ]
        
        # 当前关键词
        self.keyword = ""
        
        # 如果文件不存在，创建文件并写入表头
        if not os.path.exists(self.csv_file):
            self._create_csv()
            
    def set_keyword(self, keyword):
        """设置当前搜索关键词，但不改变文件名"""
        self.keyword = keyword
        # 不再修改CSV文件名，使用固定名称
        # 确保文件存在
        if not os.path.exists(self.csv_file):
            self._create_csv()

    def _create_csv(self):
        """创建CSV文件并写入表头"""
        with open(self.csv_file, 'w', encoding='utf-8-sig', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(self.headers)
            print(f"创建数据文件：{self.csv_file}")

    def collect_article_data(self, page):
        """采集文章数据"""
        try:
            # 等待页面加载完成
            time.sleep(3)
            
            # 获取文章标题 (处理没有标题的情况)
            try:
                title_element = page.ele('xpath://*[@id="detail-title"]')
                
                # 调试信息
                print(f"标题元素: {title_element}")
                if title_element:
                    raw_title = title_element.text
                    print(f"文章标题: '{raw_title}'")
                    title = raw_title.strip() if raw_title else ""
                    # 不再过滤表情符号，保留原始文本
                else:
                    title = ""
                    
            except Exception as e:
                title = ""  # 如果没有标题，设置为空字符串
                print(f"注意：获取笔记标题时出错: {e}")
            
            # 获取作者
            try:
                # 尝试多种XPath定位作者
                author_element = page.ele('xpath://*[@id="noteContainer"]/div[4]/div[1]/div/div[1]/a[2]/span')
                if not author_element:
                    author_element = page.ele('xpath://span[@class="user-name"]')
                if not author_element:
                    author_element = page.ele('xpath://a[contains(@class, "author")]//span')
                
                print(f"作者元素: {author_element}")
                if author_element:
                    raw_author = author_element.text
                    print(f"作者: '{raw_author}'")
                    author = raw_author.strip() if raw_author else ""
                    # 不再过滤表情符号，保留原始文本
                else:
                    author = ""
                    
            except Exception as e:
                author = ""
                print(f"注意：获取作者信息时出错: {e}")
            
            # 获取文章内容
            try:
                content_element = page.ele('xpath://*[@id="detail-desc"]')
                if content_element:
                    raw_content = content_element.text
                    content = raw_content.strip() if raw_content else ""
                    # 不再过滤表情符号，保留原始文本
                else:
                    content = ""
            except Exception:
                content = ""
                print("注意：无法获取文章内容")
            
            # 获取编辑时间
            try:
                edit_time = page.ele('xpath://*[@id="noteContainer"]/div[4]/div[2]/div[1]/div[3]/span[1]').text.strip()
            except Exception:
                edit_time = ""
                print("注意：无法获取编辑时间")
            
            # 获取点赞量
            try:
                likes = page.ele('xpath://*[@id="noteContainer"]/div[4]/div[3]/div/div/div[1]/div[2]/div/div[1]/span[1]/span[2]').text.strip()
            except Exception:
                likes = "0"
                print("注意：无法获取点赞量")
            
            # 获取收藏量
            try:
                collects = page.ele('xpath://*[@id="note-page-collect-board-guide"]/span').text.strip()
            except Exception:
                collects = "0"
                print("注意：无法获取收藏量")
            
            # 获取评论量
            try:
                comments = page.ele('xpath://*[@id="noteContainer"]/div[4]/div[3]/div/div/div[1]/div[2]/div/div[1]/span[3]/span').text.strip()
            except Exception:
                comments = "0"
                print("注意：无法获取评论量")
            
            # 获取当前URL
            url = page.url
            
            # 获取采集时间
            collect_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            
            # 保存数据前打印调试信息
            print(f"即将保存数据: 标题='{title}', 作者='{author}'")
            
            # 检查是否已存在相同URL的数据
            if not self._is_url_exists(url):
                # 保存数据（添加关键词）
                self._save_to_csv([title, author, content, edit_time, likes, collects, comments, collect_time, url, self.keyword])
                print(f"已采集文章：{title if title else '无标题笔记'}")
                return True
            else:
                print(f"文章已存在，跳过：{title if title else '无标题笔记'}")
                return False
            
        except Exception as e:
            print(f"采集数据时出错: {e}")
            return False

    def _is_url_exists(self, url):
        """检查URL是否已存在于CSV文件中"""
        try:
            if not os.path.exists(self.csv_file):
                return False
                
            with open(self.csv_file, 'r', encoding='utf-8-sig') as f:
                reader = csv.reader(f)
                next(reader)  # 跳过表头
                urls = [row[8] for row in reader if len(row) > 8]  # URL在第9列
                return url in urls
        except Exception:
            return False

    def _save_to_csv(self, data):
        """保存数据到CSV文件"""
        with open(self.csv_file, 'a', encoding='utf-8-sig', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(data)
            
    def _filter_emoji(self, text):
        """过滤掉文本中的表情符号，保留中文和其他字符"""
        if not text:
            return ""
        
        # 使用正则表达式过滤表情符号，但保留中文字符
        import re
        # 这个模式匹配大多数Unicode表情符号，但不会匹配中文字符
        emoji_pattern = re.compile(
            "["
            "\U0001F600-\U0001F64F"  # 表情符号
            "\U0001F300-\U0001F5FF"  # 符号和象形文字
            "\U0001F680-\U0001F6FF"  # 交通和地图符号
            "\U0001F700-\U0001F77F"  # 字母符号
            "\U0001F780-\U0001F7FF"  # 几何形状
            "\U0001F800-\U0001F8FF"  # 补充箭头
            "\U0001F900-\U0001F9FF"  # 补充符号和象形文字
            "\U0001FA00-\U0001FA6F"  # 扩展符号
            "\U0001FA70-\U0001FAFF"  # 表情符号扩展
            "\U00002702-\U000027B0"  # 装饰符号
            "\U000024C2-\U0001F251" 
            "]+", flags=re.UNICODE
        )
        
        filtered_text = emoji_pattern.sub(r'', text)
        
        # 调试输出
        if filtered_text != text:
            print(f"原始文本: '{text}'")
            print(f"过滤后文本: '{filtered_text}'")
            
        return filtered_text 